import re
import json
import pandas as pd
from typing import Dict, List, Any, Optional, Union
from loguru import logger
import sqlparse
from datetime import datetime, date

class CustomJSONEncoder(json.JSONEncoder):
    """自定义JSON编码器，处理特殊数据类型"""
    def default(self, obj):
        # 处理日期时间类型
        if isinstance(obj, (datetime, date)):
            return obj.isoformat()
        # 处理pandas的Timestamp类型
        if hasattr(obj, 'isoformat'):
            return obj.isoformat()
        # 处理其他特殊类型
        try:
            return str(obj)
        except:
            return super().default(obj)

def format_results(results: List[Dict[str, Any]], format_type: str = "text") -> str:
    """
    格式化查询结果为不同的输出格式
    
    Args:
        results: 查询结果列表
        format_type: 输出格式类型，支持 'text', 'json', 'html', 'csv'
        
    Returns:
        格式化后的结果字符串
    """
    if not results:
        return "没有结果"
    
    # 在转换为DataFrame前，处理特殊类型
    processed_results = []
    for row in results:
        processed_row = {}
        for key, value in row.items():
            # 处理日期时间类型
            if isinstance(value, (datetime, date)):
                processed_row[key] = value.isoformat()
            # 处理其他可能导致问题的类型
            elif value is not None and not isinstance(value, (str, int, float, bool)):
                try:
                    # 尝试转换为字符串
                    processed_row[key] = str(value)
                except:
                    processed_row[key] = "无法显示的值类型"
            else:
                processed_row[key] = value
        processed_results.append(processed_row)
    
    # 将结果转换为DataFrame
    try:
        df = pd.DataFrame(processed_results)
    except Exception as e:
        logger.error(f"转换结果为DataFrame失败: {str(e)}")
        # 返回原始结果的字符串表示
        return f"处理结果时出错: {str(e)}\n原始结果: {str(processed_results)[:1000]}..."
    
    if format_type.lower() == "json":
        try:
            return json.dumps(processed_results, ensure_ascii=False, indent=2, cls=CustomJSONEncoder)
        except Exception as e:
            logger.error(f"转换为JSON失败: {str(e)}")
            return f"转换为JSON失败: {str(e)}"
    
    elif format_type.lower() == "html":
        try:
            return df.to_html(index=False)
        except Exception as e:
            logger.error(f"转换为HTML失败: {str(e)}")
            return f"转换为HTML失败: {str(e)}"
    
    elif format_type.lower() == "csv":
        try:
            return df.to_csv(index=False)
        except Exception as e:
            logger.error(f"转换为CSV失败: {str(e)}")
            return f"转换为CSV失败: {str(e)}"
    
    else:  # 默认为文本格式
        try:
            # 获取列宽度
            col_widths = {}
            for col in df.columns:
                # 计算列标题的宽度和最大数据宽度
                col_widths[col] = max(
                    len(str(col)),
                    df[col].astype(str).str.len().max() if not df.empty else 0
                )
            
            # 构建表头
            header = " | ".join(f"{col:{col_widths[col]}s}" for col in df.columns)
            separator = "-+-".join("-" * col_widths[col] for col in df.columns)
            
            # 构建数据行
            rows = []
            for _, row in df.iterrows():
                formatted_row = " | ".join(
                    f"{str(row[col]):{col_widths[col]}s}" for col in df.columns
                )
                rows.append(formatted_row)
            
            # 组合结果
            result = f"{header}\n{separator}\n" + "\n".join(rows)
            return result
        except Exception as e:
            logger.error(f"转换为文本格式失败: {str(e)}")
            return f"转换为文本格式失败: {str(e)}"

def validate_sql(query: str) -> Dict[str, Any]:
    """
    验证SQL查询的安全性和有效性
    
    Args:
        query: SQL查询字符串
        
    Returns:
        包含验证结果和可能的警告/错误的字典
    """
    result = {
        "is_valid": True,
        "warnings": [],
        "errors": [],
        "query_type": None,
        "tables": []
    }
    
    # 使用sqlparse解析查询
    try:
        statements = sqlparse.parse(query)
        if not statements:
            result["is_valid"] = False
            result["errors"].append("无法解析SQL查询")
            return result
        
        # 获取查询类型
        stmt = statements[0]
        stmt_type = stmt.get_type().lower() if stmt.get_type() else None
        result["query_type"] = stmt_type
        
        # 检查查询类型
        if stmt_type in ["delete", "drop", "truncate"]:
            result["warnings"].append(f"请谨慎使用 {stmt_type.upper()} 操作，这可能会删除数据")
        
        if stmt_type in ["update", "insert"]:
            result["warnings"].append(f"{stmt_type.upper()} 操作将修改数据库内容")
        
        # 提取表名
        table_pattern = r"(?:FROM|JOIN|UPDATE|INTO)\s+[\"]?([a-zA-Z0-9_\.]+)[\"]?"
        tables = re.findall(table_pattern, query, re.IGNORECASE)
        result["tables"] = list(set(tables))  # 去重
        
        # 检查常见的安全问题
        if ";" in query[:-1]:  # 查询中间有分号，可能是SQL注入尝试
            result["warnings"].append("查询中含有多个语句，可能存在安全风险")
        
        if re.search(r"--", query):
            result["warnings"].append("查询中含有注释，请确保这是有意为之")
        
        if re.search(r"/\*.*?\*/", query, re.DOTALL):
            result["warnings"].append("查询中含有块注释，请确保这是有意为之")
        
        # 检查表是否为空
        if not result["tables"] and stmt_type in ["select", "update", "delete"]:
            result["warnings"].append("未能识别查询中的表名")
        
        return result
    
    except Exception as e:
        result["is_valid"] = False
        result["errors"].append(f"验证SQL时出错: {str(e)}")
        logger.error(f"验证SQL失败: {str(e)}")
        return result 